import networkx as nx
from typing import Dict, Tuple, List, Set
from collections import defaultdict

import self

from topo2 import build_hybrid_topology_shared_tor
from flow2 import TrafficGenerator


def allocate_switches_with_es(
        G: nx.Graph,
        T: List[Tuple[str, str]],
        D: Dict[Tuple[str, str], float],
        tor_switches: Set[str],
        optical_switches: Set[str],
        k_paths: int = 3
) -> Tuple[Dict, Dict]:
    """
    基于流量需求的ToR交换机和光交换机路径分配算法

    参数：
    G: 网络拓扑图（节点需标记类型，边需包含capacity属性）
    T: 通信对列表 [(t1, t2), ...]
    D: 流量需求字典 {(t1, t2): demand}
    tor_switches: ToR交换机集合
    optical_switches: 光交换机集合
    k_paths: 最大路径查找数量

    返回：
    S: 交换机分配方案 {(t1, t2): (tor1, tor2)}
    P: 路径分配方案 {(t1, t2): path}
    """

    # 初始化数据结构
    S = {}
    P = {}

    # 初始化链路剩余容量（使用有序边元组）
    remaining_capacity = {}
    for u, v in G.edges():
        edge = tuple(sorted((u, v)))
        remaining_capacity[edge] = G[u][v]['capacity']

    # 初始化光交换机占用状态
    switch_occupied = {node: False for node in optical_switches}

    # 按流量需求降序排序通信对
    sorted_pairs = sorted(T, key=lambda x: D[x], reverse=True)

    # 处理每个通信对
    for t1, t2 in sorted_pairs:
        demand = D[(t1, t2)]

        # 获取连接的ToR交换机
        tor1 = [n for n in G.neighbors(t1) if n in tor_switches]
        tor2 = [n for n in G.neighbors(t2) if n in tor_switches]

        if not tor1 or not tor2:
            S[(t1, t2)] = None
            P[(t1, t2)] = None
            continue

        # 查找所有可行路径
        all_paths = []
        for src in tor1:
            for dst in tor2:
                try:
                    # 使用Yen's算法查找前k条路径
                    paths = nx.shortest_simple_paths(G, src, dst, weight='hops')
                    all_paths.extend([(src, dst, p) for p, _ in zip(paths, range(k_paths))])
                except nx.NetworkXNoPath:
                    continue

        # 过滤可行路径
        feasible_paths = []
        for src, dst, path in all_paths:
            if is_path_feasible(path, optical_switches, switch_occupied, remaining_capacity, demand):
                feasible_paths.append((src, dst, path))

        if not feasible_paths:
            S[(t1, t2)] = None
            P[(t1, t2)] = None
            continue

        # 选择最优路径
        best_path = select_optimal_path(feasible_paths, remaining_capacity)
        src, dst, path = best_path

        # 记录分配方案
        S[(t1, t2)] = (src, dst)
        P[(t1, t2)] = path

        # 更新网络资源
        update_resources(path, remaining_capacity, demand, optical_switches, switch_occupied)

    return S, P


def is_path_feasible(
        path: List[str],
        optical_switches: Set[str],
        switch_occupied: Dict[str, bool],
        remaining_capacity: Dict[Tuple[str, str], float],
        demand: float
) -> bool:
    """检查路径是否可行"""
    # 检查光交换机占用状态
    for node in path:
        if node in optical_switches and switch_occupied[node]:
            return False

    # 检查链路容量
    for i in range(len(path) - 1):
        u, v = sorted((path[i], path[i + 1]))
        edge = (u, v)
        if remaining_capacity[edge] < demand:
            return False

    return True


def select_optimal_path(
        feasible_paths: List[Tuple[str, str, List[str]]],
        remaining_capacity: Dict[Tuple[str, str], float]
) -> Tuple[str, str, List[str]]:
    """选择最优路径（最大瓶颈容量）"""
    max_capacity = -1
    best_path = None

    for src, dst, path in feasible_paths:
        min_cap = float('inf')
        for i in range(len(path) - 1):
            u, v = sorted((path[i], path[i + 1]))
            min_cap = min(min_cap, remaining_capacity[(u, v)])

        if min_cap > max_capacity:
            max_capacity = min_cap
            best_path = (src, dst, path)

    return best_path


def update_resources(
        path: List[str],
        remaining_capacity: Dict[Tuple[str, str], float],
        demand: float,
        optical_switches: Set[str],
        switch_occupied: Dict[str, bool]
) -> None:
    """更新网络资源状态"""
    # 更新链路容量
    for i in range(len(path) - 1):
        u, v = sorted((path[i], path[i + 1]))
        remaining_capacity[(u, v)] -= demand

    # 标记光交换机
    for node in path:
        if node in optical_switches:
            switch_occupied[node] = True


# 测试用例
if __name__ == "__main__":
    # 创建测试拓扑
    G = build_hybrid_topology_shared_tor()  # 从函数中获取拓扑
    print("---------------------------------------------------------")
    # 检查图的节点
    print("图的节点G.nodes:", G.nodes()) # 获取到的是所有的节点，包括交换机和服务器
    generator = TrafficGenerator(num_nodes=16) # 从函数中获取流量生成器，输入的16是ToR的数量
    D = generator.generate_all_to_all() # 生成所有节点对之间的流量（ToR节点对之间的流量）

    # 将 numpy.ndarray 转换为字典
    D_dict = {(i, j): D[i][j] for i in range(D.shape[0]) for j in range(D.shape[1])}
    print("D_dict",D_dict)
    # # 获取 ToR 交换机和光交换机集合
    tor_switches = {node for node, data in G.nodes(data=True) if data.get('type') == 'tor'}
    print("tor_switches",tor_switches)
    optical_switches = {node for node, data in G.nodes(data=True) if data.get('type') == 'optical_switch'}
    print("光交换机节点为：",optical_switches)
    print("--------------------------------------------------------------------------------")
    print("D_dict.items()",D_dict.items())
    # 生成通信对列表
    T = []
    for pair, demand in D_dict.items():
        # print("pair:",pair)
        # print("demand:",demand)
        if demand > 0:
            T.append(pair)
    for pair in T:
        print("pair:",pair) # 输出的通信对的格式为（1,2），表示ToR1和ToR2
    #
    # # 检查通信对节点是否在图中
    # for pair in T:
    #     for node in pair:
    #         if node not in G.nodes():
    #             print(f"节点 {node} 不在图 G 中")
    #
    # S, P = allocate_switches_with_es(G, T, D_dict, tor_switches, optical_switches)
    #
    # # 输出结果
    # print("交换机分配方案:")
    # for pair, assignment in S.items():
    #     print(f"{pair} -> {assignment}")
    #
    # print("\n路径分配方案:")
    # for pair, path in P.items():
    #     print(f"{pair} -> {path}")
    # 假设 D 是 generate_all_to_all() 返回的字典

    # 将键从 (i, j) 转换为 ("ToR{i}", "ToR{j}")
    D_str_keys = {(f"ToR{i}", f"ToR{j}"): value for (i, j), value in D_dict.items()}
    # print("D_str_keys:",D_str_keys)

    # 将值从 numpy.ndarray 转换为浮点数
    D_float_values = {k: float(v) for k, v in D_str_keys.items()}

    # 生成通信对列表
    T = [pair for pair, demand in D_float_values.items() if demand > 0]
    print("生成通信对列表",T)
    print("T[32]:",T[32])

    # 检查通信对节点是否在图中
    # for pair in T:
    #     for node in pair:
    #         if node not in G.nodes():
    #             print(f"节点 {node} 不在图 G 中")
    #         else:
    #             print(f"节点 {node} 在图 G 中")

    # 调用分配算法
    # S, P = allocate_switches_with_es(G, T, D_float_values, tor_switches, optical_switches)
    #
    # # 输出结果
    # print("交换机分配方案:")
    # for pair, assignment in S.items():
    #     print(f"{pair} -> {assignment}")
    #
    # print("\n路径分配方案:")
    # for pair, path in P.items():
    #     print(f"{pair} -> {path}")

    # 节点之间存在路径
    print("=====================================================")
    pair1=T[32]
    src=pair1[0]
    dst=pair1[1]
    # try:
    #     paths = nx.shortest_simple_paths(G, src, dst, weight='hops')
    #     for path in paths:
    #         print(f"从 {src} 到 {dst} 的路径: {path}")
    # except nx.NetworkXNoPath:
    #     print(f"从 {src} 到 {dst} 没有路径")
